airbyte/exceptions.py (1)

415-421: Optional: include primary_key_value in safe logging context?

AirbyteRecordNotFoundError looks consistent with the surrounding connector errors. Would it be useful to also surface primary_key_value from this exception in PyAirbyteError.safe_logging_dict (similar to stream_name), so logs have enough context when a lookup fails, assuming the PK isn’t considered sensitive in your threat model, wdyt?

tests/unit_tests/test_get_record.py (2)

24-57: Stabilize expectations around fetch_record “happy path”?

In test_declarative_executor_fetch_record_stream_validation, the “valid_stream_and_pk” case currently asserts that fetch_record raises one of (NotImplementedError, AttributeError, KeyError) (Line [56]), i.e. any internal failure after stream lookup. As the CDK integration evolves and ModelToComponentFactory wiring stabilizes, this path may eventually succeed, which would flip this test from “green” to “red” even though behavior improved.

Would it be more future‑proof to explicitly assert only the stream‑existence contract here—for example by mocking out ModelToComponentFactory.create_component and SimpleRetriever so that the test simply verifies that no AirbyteStreamNotFoundError is raised for an existing stream, and leaves the downstream behavior to more targeted tests, wdyt?

---

144-149: Use typing.Any instead of bare any in type hints?

Both test_source_normalize_and_validate_pk_value and test_source_get_record_accepts_various_pk_formats annotate pk_value as any. With from __future__ import annotations, that ends up as the string "any", which type checkers typically won’t recognize as Any.

Would you consider importing Any from typing and updating these annotations to pk_value: Any for clarity and better static‑analysis support, wdyt?

Also applies to: 217-217

airbyte/_executors/declarative.py (1)

18-25: Double‑check CDK integration and edge cases in fetch_record

The high‑level flow in fetch_record looks aligned with the intended design, but a few CDK integration details and edge cases might be worth validating before relying on this in production:

1. Factory usage and manifest shape (Lines [192]-[207])
   Here you call ModelToComponentFactory.create_component with model_type=type(retriever_config) where retriever_config is pulled directly from the manifest dict. My understanding is that CDK factories typically expect a Pydantic model instance (or specific declarative model types) as component_definition, not a raw dict. Depending on the exact CDK version/contract, this could either work by accident or fail at runtime and surface as the NotImplementedError you wrap around it.
   Would it make sense to (a) confirm what create_component expects in the current CDK, and/or (b) add a small integration test using a real declarative manifest to ensure this wiring actually produces a SimpleRetriever, wdyt?

2. Schema/record selector wiring (Lines [262]-[273])
   records_schema is currently derived from stream_config.get("schema_loader", {}), but in many declarative manifests schema_loader is a configuration for a loader rather than the resolved JSON schema. Are you sure record_selector.select_records expects the loader config here, rather than the actual JSON schema (or even None), and if not, should we either:

   • resolve the schema via the loader, or
   • explicitly pass {} and rely on selectors that don’t need schema,
     wdyt?

3. Empty records but non‑dict JSON responses (Lines [275]-[282])
   When records is empty, you only special‑case non‑empty dict JSON bodies. If a connector returns [{"id": "123", ...}] (a single‑element list) for a detail endpoint, this will currently raise AirbyteRecordNotFoundError even though the response is there. Would you consider handling the “single‑element list” case by returning that single element (and perhaps logging a warning), or is it intentional to only support dict‑shaped bodies for now, wdyt?

4. HTTP status / 404 handling
   Since you call send_request directly, a 404 might be represented as either:

   • an exception from the requester, or
   • a normal response object with status 404 but no body/records.
     If it’s the latter, would you want to explicitly check response.status_code (when available) and map a 404 to AirbyteRecordNotFoundError for clearer semantics, wdyt?

5. Path construction & URL encoding (Line [225])
   fetch_path = f"{base_path}/{primary_key_value}".lstrip("/") works well for numeric/simple IDs, but if a primary key ever contains reserved URL characters (e.g., spaces, slashes), the request path could be malformed. Would it be safer to quote the primary_key_value (e.g., via urllib.parse.quote) before appending, or at least leave a comment/TODO here to revisit once you see real‑world connectors, wdyt?

Given your PR description already calls out that this isn’t yet exercised against real connectors, adding a small note or follow‑up task around these points could help ensure the API behaves as expected once wired into actual YAML manifests, wdyt?

Also applies to: 150-297

airbyte/sources/base.py (3)

605-621: Clarify interaction between primary‑key overrides and nested PK shapes

_get_stream_primary_key nicely flattens the “nested list” Airbyte PK structure (e.g., [[\"id\"], [\"org_id\"]] → ["id", "org_id"]). One subtle corner case is how this interacts with overrides coming from set_primary_key:

• set_primary_key("stream", ["id", "org_id"]) stores ["id", "org_id"] in _primary_key_overrides.
• get_configured_catalog then wraps that as [self._primary_key_overrides[...]], yielding primary_key=[[\"id\", \"org_id\"]] for the configured stream.
• _get_stream_primary_key sees pk[0] as a list and flattens via [field[0] if isinstance(field, list) else field for field in pk], which would return only "id" in this specific override shape.

Would it be worth either:

• clarifying in docs that set_primary_key should be used only for single‑field PKs for now, or
• adjusting the flattening logic/tests to cover the override case so composite PK overrides don’t get silently truncated, wdyt?

---

628-675: Tighten validation of pk_value (e.g., None) before coercing to string?

_normalize_and_validate_pk_value does a good job of enforcing single‑field PKs and validating dict input against the expected field. Right now, any non‑dict value—including None or empty strings—is accepted and coerced via str(pk_value) (Line [675]), which could lead to requests like /users/None or /users/ if a caller accidentally passes a falsy/placeholder value.

Would you consider:

• explicitly rejecting pk_value is None (and perhaps pk_value == "") with a PyAirbyteInputError, and
• maybe including the stream name and PK field in the error message to aid debugging,

so consumers get a clear input‑validation failure instead of a surprising HTTP path, wdyt?

---

677-714: Confirm get_record is the only entrypoint that should call _normalize_and_validate_pk_value

get_record is currently the single public entrypoint using _normalize_and_validate_pk_value, and it correctly enforces DeclarativeExecutor before delegating. Given that _normalize_and_validate_pk_value is fairly generic, do you want to keep it private and only reachable via get_record, or would you prefer to document it (or expose a wrapper) for advanced users building their own request logic on top of Source?

If you intend it to remain strictly internal, maybe a brief comment above _normalize_and_validate_pk_value noting that it’s “used only by get_record for now” could help future maintainers avoid calling it in incompatible contexts, wdyt?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: Suraj-Vishwakarma70
Repo: airbytehq/PyAirbyte PR: 285
File: airbyte/sources/base.py:0-0
Timestamp: 2024-10-08T15:34:31.026Z
Learning: Ensure consistent naming for attributes in the `Source` class in `airbyte/sources/base.py`, such as renaming `_to_be_selected_stream` to `_to_be_selected_streams`.

airbyte/_executors/declarative.py (1)

325-342: Tighten the JSON fallback and replace the bare except Exception: pass

The try/except Exception: pass around response.json() on Lines [327]-[331] both repeats the earlier “empty except” concern and makes it harder to notice unexpected decode failures. Would you consider:

• Switching to with contextlib.suppress(Exception): around the response.json() call (as you already do for schema_loader.get_json_schema()), so the intent to ignore parse errors is explicit without an empty except; and
• Optionally also handling the case where response.json() returns a non-empty list (e.g. a one-element [{...}]) by treating a single dict element as the record, to cover APIs that return arrays even for singleton lookups?

That would preserve the current not-found semantics while addressing the lint concern and slightly widening compatibility, wdyt?

-        if not records:
-            try:
-                response_json = response.json()
-                if isinstance(response_json, dict) and response_json:
-                    return response_json
-            except Exception:
-                pass
+        if not records:
+            with contextlib.suppress(Exception):
+                response_json = response.json()
+                if isinstance(response_json, dict) and response_json:
+                    return response_json
+                if (
+                    isinstance(response_json, list)
+                    and len(response_json) == 1
+                    and isinstance(response_json[0], dict)
+                ):
+                    return response_json[0]

airbyte/_executors/declarative.py (1)

45-89: Consider iterating through wrapper chains in _unwrap_to_declarative_stream

Right now the helper only peeks one level through the known wrapper attributes and branch attrs (Lines [64]-[83]), so a future wrapper-of-wrapper that still ultimately exposes a retriever via the same attributes would end up raising NotImplementedError even though it could be unwrapped. Would you consider turning this into an iterative loop (tracking visited objects and repeatedly following wrapper_attrs / branch_attr until a retriever is found or a depth limit is reached) to make it more robust to additional wrapper layers while keeping the same error behavior when no retriever exists, wdyt?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

airbyte/_executors/declarative.py (1)

321-327: Address the silent exception suppression with logging.

The empty except block at lines 326-327 silently suppresses exceptions when attempting to parse response.json(). This was flagged in previous review comments as needing an explanatory comment or logging.

While the current behavior (try to parse JSON, fall through to error if it fails) is reasonable, adding logging would help with debugging when unexpected responses occur.

Would you consider adding logging here as suggested in the previous review? For example:

+import logging
+
+logger = logging.getLogger(__name__)

# ... in fetch_record method ...

         if not records:
             try:
                 response_json = response.json()
                 if isinstance(response_json, dict) and response_json:
                     return response_json
-            except Exception:
-                pass
+            except Exception as e:
+                logger.debug(
+                    "Failed to parse response as JSON when fetching record for stream '%s', "
+                    "primary key '%s': %s",
+                    stream_name,
+                    primary_key_value,
+                    str(e),
+                )

This provides visibility into parsing failures while maintaining the existing fallthrough behavior.

airbyte/sources/base.py (2)

605-626: Consider adding type validation for primary key elements, wdyt?

At line 620, when pk is a flat list (not nested), we cast it to list[str] with a type ignore comment. However, there's no explicit validation that the elements are actually strings. If the catalog contains unexpected types, this could propagate invalid data downstream.

Would it be safer to explicitly validate and convert elements to strings? For example:

                 if isinstance(pk, list) and len(pk) > 0:
                     if isinstance(pk[0], list):
                         return [field[0] if isinstance(field, list) else field for field in pk]
-                    return list(pk)  # type: ignore[arg-type]
+                    return [str(field) for field in pk]

This would ensure type safety and make the conversion explicit. Alternatively, if we trust the catalog format completely, a comment explaining why the type ignore is safe would help future maintainers.

---

628-675: Consider validating against None and empty primary key values, wdyt?

The method converts pk_value to a string without checking for None or empty values. If a user accidentally passes None, it would be converted to the string "None", which is probably not the intended primary key value. Similarly, empty strings or whitespace-only strings might warrant validation.

Would it make sense to add validation like this after line 675?

         return str(pk_value[provided_key])
 
     return str(pk_value)
+
+def _normalize_and_validate_pk_value(
+    self,
+    stream_name: str,
+    pk_value: Any,  # noqa: ANN401
+) -> str:
+    """Normalize and validate a primary key value."""
+    # ... existing logic ...
+    
+    result = str(pk_value)  # or the dict extraction above
+    
+    if not result or result == "None":
+        raise exc.PyAirbyteInputError(
+            message=f"Primary key value cannot be None or empty for stream '{stream_name}'.",
+            input_value=str(pk_value),
+        )
+    
+    return result

This would catch common user errors early with a clear message.

airbyte/_executors/declarative.py (3)

45-89: Consider validating that retriever is not None, wdyt?

The unwrapping logic checks hasattr(unwrapped, "retriever") at lines 76 and 82, but doesn't verify that the retriever attribute itself is not None. If a stream wrapper has a retriever attribute set to None, this function would return the wrapper, and the caller at line 249 would attempt to access it without a null check.

Would it be safer to add a not-None check? For example:

         if hasattr(stream, attr_name):
             unwrapped = getattr(stream, attr_name)
-            if unwrapped is not None and hasattr(unwrapped, "retriever"):
+            if unwrapped is not None and hasattr(unwrapped, "retriever") and unwrapped.retriever is not None:
                 return unwrapped

Similarly for the branch attributes at lines 81-83. This would make the function more robust against unexpected CDK states.

---

257-267: Consider adding validation for path construction edge cases, wdyt?

The path construction at lines 264-267 assumes that base_path is well-formed and simply appends the primary key value. However, there are potential edge cases:

1. What if base_path ends with multiple slashes (e.g., "/users//")?
2. What if base_path already contains a path segment that looks like a placeholder?
3. What if primary_key_value contains URL-unsafe characters?

Would it be worth adding some defensive checks or URL encoding? For example:

from urllib.parse import quote

# After line 262
if base_path:
    # Normalize the base path and encode the PK value
    normalized_base = base_path.rstrip('/')
    encoded_pk = quote(str(primary_key_value), safe='')
    fetch_path = f"{normalized_base}/{encoded_pk}"
else:
    fetch_path = quote(str(primary_key_value), safe='')

This would handle special characters in PK values (like spaces, slashes, etc.) more robustly. Though if this isn't a concern for the initial implementation, it could be deferred.

---

195-343: Consider breaking down fetch_record into smaller helper methods, wdyt?

The fetch_record method is quite long (149 lines) and handles multiple responsibilities: stream lookup, retriever validation, path construction, HTTP request, schema retrieval, record selection, and response parsing. This complexity is reflected in the linter suppressions (PLR0914, PLR0912, PLR0915).

While the current implementation is functional, would it be beneficial to extract some of these concerns into smaller, testable helper methods? For example:

def _find_stream(self, stream_name: str, streams: list) -> object:
    """Find and return the target stream by name."""
    # Lines 220-245
    
def _build_fetch_path(self, retriever: SimpleRetriever, primary_key_value: str) -> str:
    """Construct the fetch path by appending PK to base path."""
    # Lines 257-267
    
def _send_fetch_request(self, retriever: SimpleRetriever, fetch_path: str, ...) -> Response:
    """Send the HTTP request to fetch the record."""
    # Lines 269-290
    
def _extract_record_from_response(self, response, retriever, stream_name, primary_key_value, ...) -> dict:
    """Extract and return the record from the HTTP response."""
    # Lines 304-343

This would make the code more maintainable, easier to test individual components, and reduce cognitive complexity. However, if you prefer to keep it as-is for the initial implementation, that's also reasonable.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: Suraj-Vishwakarma70
Repo: airbytehq/PyAirbyte PR: 285
File: airbyte/sources/base.py:0-0
Timestamp: 2024-10-08T15:34:31.026Z
Learning: Ensure consistent naming for attributes in the `Source` class in `airbyte/sources/base.py`, such as renaming `_to_be_selected_stream` to `_to_be_selected_streams`.

tests/unit_tests/test_get_record.py (1)

17-90: Consider adding tests for additional DeclarativeExecutor edge cases?

Looking at the implementation in airbyte/_executors/declarative.py, a few edge cases could strengthen coverage:

1. Non-SimpleRetriever case: When a stream uses a different retriever type (should raise NotImplementedError)
2. Response is None: When send_request() returns None (should raise AirbyteRecordNotFoundError)
3. Empty records with fallback: When select_records() returns empty but response.json() has data
4. Stream unwrapping: When streams are wrapped and need unwrapping to find the declarative stream

These would catch regressions in the error-handling paths. Given that the PR notes this hasn't been tested with real connectors yet, these additional tests might catch issues before integration testing, wdyt?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

airbyte/sources/base.py (2)

606-638: Consider simplifying by using CatalogProvider directly, wdyt?

The current implementation creates a temporary catalog (lines 623-630) just to call CatalogProvider.get_primary_keys(). Since self.configured_catalog already has properly formatted primary keys, could we simplify this to:

def _get_stream_primary_key(self, stream_name: str) -> list[str]:
    """Get the primary key for a stream."""
    catalog_provider = CatalogProvider(self.configured_catalog)
    return catalog_provider.get_primary_keys(stream_name)

This would:

• Eliminate the manual catalog iteration and temporary object creation
• Reuse the validation logic already in CatalogProvider
• Let CatalogProvider.get_configured_stream_info() handle the stream-not-found case consistently

The defensive check for flat vs nested PK format (lines 620-621) shouldn't be needed since get_configured_catalog() (lines 452-456) already ensures primary keys are in the correct nested format.

---

719-740: Minor optimization: pk_value_str computed twice in fallback path

When a DeclarativeExecutor falls back to scanning (lines 726-729), pk_value_str is computed twice:

• Line 720: For the direct fetch attempt
• Line 740: For the scanning fallback

You could optimize by computing once before line 719:

+        pk_value_str = self._normalize_and_validate_pk_value(stream_name, pk_value)
         if isinstance(self.executor, DeclarativeExecutor):
-            pk_value_str = self._normalize_and_validate_pk_value(stream_name, pk_value)
             try:
                 return self.executor.fetch_record(

Then remove line 740 since pk_value_str is already available. This avoids redundant validation when falling back to scanning, wdyt?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: Suraj-Vishwakarma70
Repo: airbytehq/PyAirbyte PR: 285
File: airbyte/sources/base.py:0-0
Timestamp: 2024-10-08T15:34:31.026Z
Learning: Ensure consistent naming for attributes in the `Source` class in `airbyte/sources/base.py`, such as renaming `_to_be_selected_stream` to `_to_be_selected_streams`.

tests/unit_tests/test_get_record.py (4)

19-92: Stream validation test is solid; consider also asserting on the constructed fetch path?

The parametrized coverage for existing vs missing streams looks good and the mocking of declarative_source.streams / SimpleRetriever aligns with the implementation. Would it be worth adding an assertion that send_request is called with the expected path (including the /primary_key_value suffix) so regressions in path construction are caught here as well, wdyt?

---

94-150: _get_stream_primary_key test currently leans on CatalogProvider internals

These tests indirectly depend on CatalogProvider behavior via _get_stream_primary_key, which is nice for integration coverage but could become brittle if CDK/CatalogProvider internals change. Would you consider patching CatalogProvider.get_primary_keys (or injecting a small fake) so this test focuses purely on how Source wires primary_key shapes into the provider, while still leaving a separate integration test for the full stack, wdyt?

---

207-222: Non-declarative executor test could assert on the exact error message

The test already uses match="only supported for declarative sources", which is helpful. If the error message ever grows, would you want to tighten this a bit more (e.g., anchor with ^/$ or assert on a dedicated error subclass) so refactors don’t accidentally weaken the contract, or is the current level of coupling intentional, wdyt?

---

224-270: Executor mocks may not exercise the same isinstance path as production

In get_record tests you use Mock(spec=DeclarativeExecutor) as the executor. Given that Source.get_record currently branches on isinstance(self.executor, DeclarativeExecutor) (per airbyte/sources/base.py), a plain Mock won’t satisfy that isinstance check, so these tests might not be hitting the same code path as a real DeclarativeExecutor instance. Would it be safer to use a lightweight DeclarativeExecutor subclass test double, or create_autospec(DeclarativeExecutor, instance=True), so the isinstance branch is exercised and any future changes to the type check are properly covered, wdyt?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

airbyte/sources/base.py (1)

711-729: Critical: Field name normalization mismatch in scanning logic.

This is the bug flagged in the previous review and it's still present. Here's what's happening:

1. Line 712: pk_field comes from catalog_provider.get_primary_keys(), which returns normalized field names (e.g., "id" for a source field "ID"), as confirmed by the relevant code snippet showing LowerCaseNormalizer.normalize() is applied.

2. Line 716: get_records(stream_name) uses the default normalize_field_names=False, so records contain original field names from the source (e.g., {"ID": "123", ...}).

3. Line 729: record_data.get(pk_field) looks up the normalized name "id" in a record that has the original field name "ID", which returns None and causes the match to fail.

Example that breaks:

• Source defines primary key: "ID" (uppercase)
• After catalog normalization: pk_field = "id" (lowercase)
• Record from source: {"ID": "123", ...} (original casing preserved)
• Lookup: record_data.get("id") → None ❌
• Match fails even though the record exists

The fix is straightforward—pass normalize_field_names=True to ensure records use normalized field names:

-        for record in self.get_records(stream_name):
+        for record in self.get_records(stream_name, normalize_field_names=True):

Would you also want to add a unit test that covers a source with non-lowercase primary key field names to prevent regressions, wdyt?

airbyte/sources/base.py (2)

422-425: Consider caching the catalog_provider to avoid repeated instantiation.

The property creates a new CatalogProvider instance on each access. Looking at get_record(), it calls this property twice indirectly (via _normalize_and_validate_pk_value at line 691 and directly at line 712), creating two instances per call.

While CatalogProvider initialization is likely lightweight, would it make sense to cache the instance—perhaps as self._catalog_provider—and only recreate it when configured_catalog changes? Or is the instantiation cost negligible enough that this optimization isn't worth the added complexity, wdyt?

---

690-713: Optional: Avoid duplicate PK normalization in scanning path.

The code calls _normalize_and_validate_pk_value() twice when entering the scanning fallback:

• Line 691: First call (before attempting declarative fetch)
• Line 711: Second call (when entering scan fallback)

Since pk_value_str is already computed at line 691, you could reuse it at line 711 instead of re-normalizing. Similarly, primary_key_fields could be retrieved once.

Would it be worth caching these values to avoid the redundant work, or is the simplicity of the current structure preferred? Wdyt?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: Suraj-Vishwakarma70
Repo: airbytehq/PyAirbyte PR: 285
File: airbyte/sources/base.py:0-0
Timestamp: 2024-10-08T15:34:31.026Z
Learning: Ensure consistent naming for attributes in the `Source` class in `airbyte/sources/base.py`, such as renaming `_to_be_selected_stream` to `_to_be_selected_streams`.